package com.md.client.helper;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;

public class BodyCachingRequestWrapper extends HttpServletRequestWrapper {
  private static final Integer MAX_BYTE_LENGTH = 1024;
  private final byte[] requestBody;

  public BodyCachingRequestWrapper(final HttpServletRequest request) throws IOException {
    super(request);
    // 读取请求体内容并缓存
    InputStream inputStream = request.getInputStream();
    ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
    byte[] buffer = new byte[MAX_BYTE_LENGTH];
    int bytesRead;
    while ((bytesRead = inputStream.read(buffer)) != -1) {
      byteArrayOutputStream.write(buffer, 0, bytesRead);
    }
    requestBody = byteArrayOutputStream.toByteArray();
  }

  @Override
  public ServletInputStream getInputStream() {
    // 使用缓存的请求体内容创建新的输入流
    return new CachedServletInputStream(requestBody);
  }

  @Override
  public BufferedReader getReader() {
    // 使用缓存的请求体内容创建新的 BufferedReader
    return new BufferedReader(new InputStreamReader(getInputStream()));
  }

  // 自定义的 ServletInputStream 类
  private static class CachedServletInputStream extends ServletInputStream {
    private final ByteArrayInputStream byteArrayInputStream;

    CachedServletInputStream(final byte[] requestBody) {
      this.byteArrayInputStream = new ByteArrayInputStream(requestBody);
    }

    @Override
    public int read() {
      return byteArrayInputStream.read();
    }

    @Override
    public boolean isFinished() {
      return byteArrayInputStream.available() == 0;
    }

    @Override
    public boolean isReady() {
      return true;
    }

    @Override
    public void setReadListener(final ReadListener readListener) {
      // Optional: you can implement this to support asynchronous I/O if needed
    }
  }
}
